from mem0 import Memory
import time
from typing import Dict, Any, Optional, List
from utils.logger import emobot_logger
from config.settings import config

logger = emobot_logger.get_logger()


class MemoService:
    def __init__(self):
        self.config = {
            "llm": {
                "provider": "openai",
                "config": {
                    "model": config.MEMO_LLM_MODEL,
                    "api_key": config.MEMO_LLM_API_KEY,
                    "openai_base_url": config.MEMO_LLM_BASE_URL,
                },
            },
            "vector_store": {
                "provider": "weaviate",
                "config": {
                    "collection_name": config.MEMO_COLLECTION_NAME,
                    "cluster_url": config.MEMO_WEAVIATE_URL,
                },
            },
            "embedder": {
                "provider": "openai",
                "config": {
                    "model": config.MEMO_EMBEDDER_MODEL,
                    "openai_base_url": config.MEMO_EMBEDDER_BASE_URL,
                    "api_key": config.MEMO_EMBEDDER_API_KEY,
                },
            },
        }
        self.memory = None

    def _get_memory_client(self) -> Memory:
        if self.memory is None:
            self.memory = Memory.from_config(self.config)
        return self.memory

    def add_memory(
        self,
        messages: List[Dict[str, str]],
        user_id: str,
        metadata: Optional[Dict[str, Any]] = None,
        infer: bool = True,
    ) -> Dict[str, Any]:
        """Add conversation memory for user"""
        try:
            memory_client = self._get_memory_client()
            result = memory_client.add(
                messages, user_id=user_id, metadata=metadata or {}, infer=infer
            )
            logger.info(f"Memory added successfully for user {user_id}")
            return result
        except Exception as e:
            logger.error(f"Error adding memory for user {user_id}: {e}")
            raise

    def search_memories(
        self, query: str, user_id: str, limit: int = 10
    ) -> List[Dict[str, Any]]:
        """Search memories for user"""
        try:
            memory_client = self._get_memory_client()
            results = memory_client.search(query=query, user_id=user_id)
            logger.info(f"Found {len(results)} memories for user {user_id}")
            return results[:limit]
        except Exception as e:
            logger.error(f"Error searching memories for user {user_id}: {e}")
            return []

    def get_all_memories(self, user_id: str) -> List[Dict[str, Any]]:
        """Get all memories for user"""
        try:
            memory_client = self._get_memory_client()
            results = memory_client.get_all(user_id=user_id)
            logger.info(f"Retrieved {len(results)} memories for user {user_id}")
            return results
        except Exception as e:
            logger.error(f"Error getting all memories for user {user_id}: {e}")
            return []

    def delete_memory(self, memory_id: str, user_id: str) -> bool:
        """Delete specific memory for user"""
        try:
            memory_client = self._get_memory_client()
            memory_client.delete(memory_id=memory_id, user_id=user_id)
            logger.info(f"Memory {memory_id} deleted successfully for user {user_id}")
            return True
        except Exception as e:
            logger.error(f"Error deleting memory {memory_id} for user {user_id}: {e}")
            return False

    def delete_all_memories(self, user_id: str) -> bool:
        """Delete all memories for user"""
        try:
            memory_client = self._get_memory_client()
            memories = self.get_all_memories(user_id)
            for memory in memories:
                if "id" in memory:
                    memory_client.delete(memory_id=memory["id"], user_id=user_id)
            logger.info(f"All memories deleted successfully for user {user_id}")
            return True
        except Exception as e:
            logger.error(f"Error deleting all memories for user {user_id}: {e}")
            return False

    def update_memory_metadata(
        self, memory_id: str, user_id: str, metadata: Dict[str, Any]
    ) -> bool:
        """Update memory metadata (requires delete and re-add)"""
        try:
            memory_client = self._get_memory_client()
            memories = self.get_all_memories(user_id)

            for memory in memories:
                if memory.get("id") == memory_id:
                    messages = memory.get("messages", [])
                    new_metadata = {**memory.get("metadata", {}), **metadata}

                    memory_client.delete(memory_id=memory_id, user_id=user_id)
                    memory_client.add(
                        messages, user_id=user_id, metadata=new_metadata, infer=False
                    )

                    logger.info(
                        f"Memory {memory_id} metadata updated successfully for user {user_id}"
                    )
                    return True

            logger.warning(f"Memory {memory_id} not found for user {user_id}")
            return False
        except Exception as e:
            logger.error(
                f"Error updating memory {memory_id} metadata for user {user_id}: {e}"
            )
            return False

    def close(self):
        """Close memory client connection"""
        if self.memory is not None:
            try:
                if hasattr(self.memory, "_vector_store") and hasattr(
                    self.memory._vector_store, "_client"
                ):
                    client = self.memory._vector_store._client
                    if hasattr(client, "close"):
                        client.close()
                time.sleep(0.5)
                logger.info("Memory client closed successfully")
            except Exception as e:
                logger.error(f"Error closing memory client: {e}")
            finally:
                self.memory = None
